from .account import Account

class Strategy:
    def __init__(self, name, algo_list=None):
        self.name = name
        self.algo_list = algo_list
        self.acc = Account()

    def algo_processor(self, context):
        for algo in self.algo_list:
            if algo(context) is True: #如果algo返回True,直接不运行，本次不调仓
                return None
        return context['weights']

    def onbar(self, index, date, df_bar):

        self.acc.update_bar(date, df_bar)
        weights = self.algo_processor({'index': index, 'bar':df_bar, 'date':date,'acc':self.acc})
        if weights is not None:
            self.acc.adjust_weights(date, weights)



from engine.algos import *
class StrategyBuyHold(Strategy):
    #params={period=1...N,可选；}
    def __init__(self, name, params={}):
        algo_period = None
        if 'period' in params.keys():
            algo_period = RunPeriod(period=params['period'])
        else:
            algo_period = RunOnce()
        algo_list = [
            algo_period,
            SelectAll(),
            WeightEqually()
        ]
        super(StrategyBuyHold, self).__init__(name,algo_list)

class StrategyRolling(Strategy):

    def __init__(self, name, params={}):
        algo_list = [
            SelectBySignal(signal_buy=params['signal_buy'],signal_sell=params['signal_sell']),
            SelectTopK(K=params['K'], col=params['sort_by']),
            WeightEqually()
        ]
        super(StrategyRolling, self).__init__(name,algo_list=algo_list)

class StrategyMachineLearning(Strategy):
    def __init__(self, name, params={}):
        algo_list = [
            #SelectBySignal(signal_buy=params['signal_buy'],signal_sell=params['signal_sell']),
            SelectTopK(K=params['K'], col=params['sort_by']),
            WeightEqually()
        ]
        super(StrategyMachineLearning, self).__init__(name,algo_list=algo_list)
        self._init_model(params)


    def _init_model(self, params):
        df = params['df']
        from engine.model.gbdt import LGBModel
        m = LGBModel()
        m.fit(df)
        results = m.predict()
        df['pred_score'] = results





